package com.jrlc.springaidemo.com.jrlc.init.controller;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@RestController
@RequestMapping("ai")
public class AiController {
    @Resource
    private ChatClient chatClient;
    @Resource
    private ChatModel chatModel;
    List<Message> messages = new ArrayList<>();


    @PostConstruct
    public void init() {
        messages.add(new SystemMessage("你是一个聊天机器人，所有回复请使用中文)"));
    }

    @GetMapping(value = "/model/stream")
    public Flux<ChatResponse> modelStream(@RequestParam(value = "msg") String userMsg) {
        return chatModel.stream(
                new Prompt(
                        userMsg,
                        OpenAiChatOptions.builder()
                                .withModel("moonshot-v1-8k")
                                .withTemperature(0.8)
                                .build()
                ));
    }

    @GetMapping(value = "/model/call")
    public String modelCall(@RequestParam(value = "msg") String userMsg) {
        return chatModel.call(
                new Prompt(
                        userMsg,
                        OpenAiChatOptions.builder()
                                .withModel("moonshot-v1-8k")
                                .withTemperature(0.8)
                                .build()
                )).getResult().getOutput().getContent();
    }

    @GetMapping(value = "/chat/stream")
    public Flux<String> chatStream(@RequestParam(value = "msg") String message) {
        return chatClient.prompt().user(message).stream().content();
    }

    @GetMapping(value = "/chat/call")
    public String chatCall(@RequestParam(value = "msg") String message) {

         return chatClient.prompt().user(message).call().content();
    }
    ChatMemory chatMemory=new InMemoryChatMemory();
    Map<String,Advisor> advisors=new HashMap<>();
    @PostMapping(value = "/chat/sse")
    public SseEmitter chatCallSSe(@RequestParam(value = "msg") String message,@RequestParam(value = "uid") String uid) {
        Advisor advisor = advisors.computeIfAbsent(uid, u -> new MessageChatMemoryAdvisor(chatMemory, u, 10));
        SseEmitter emitter = new SseEmitter(100000L);
        UserMessage userMessage = new UserMessage(message);
        chatClient.prompt(new Prompt(userMessage)).advisors(advisor).stream().chatResponse().subscribe(chatResponse -> {
            String content = chatResponse.getResult().getOutput().getContent();
            if (content != null) {
                log.info("chatResponse: {}",content );
                try {
                    emitter.send(content);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        });
        return emitter;
    }
}
